# -*- coding: utf-8 -*-
# !/usr/bin/env python3

"""
@Author   : Leonis.Li
@Email    : leonis.li.ext@zeiss.com
@Cellphone:
@Version  : 1.0
@License  : Apache Licence
@Script   : utils.py
@CreateAt : 2023/10/11 22:58
@UpdateAt : 2023/10/11 22:58
"""


import base64
import cv2
import datetime
import os
import time
from typing import Union

import fitz
from paddleocr import PaddleOCR, PPStructure

from flask import current_app


# Set cache file directory
FILE_CACHE_PATH = os.environ.get("FILE_CACHE_PATH")

# Optional Target text from Source type For now
ALLOWED_EXTRACT_TEXT_FROM_DOC_OBJECTS = ["ALL", "TEXT", "TABLE", "FIGURE"]


# Set model
OCR_CH_MODEL = PaddleOCR(show_log=False,
                         use_angle_cls=True,
                         use_gpu=False,
                         lang="ch",
                         det_model_dir="/PaddleOCR/inference/ch_PP-OCRv3_det_infer/",
                         cls_model_dir="/PaddleOCR/inference/ch_ppocr_mobile_v2.0_cls_infer/",
                         rec_model_dir="/PaddleOCR/inference/ch_PP-OCRv3_rec_infer/"
                         )

OCR_EN_MODEL = PaddleOCR(show_log=False,
                         use_angle_cls=True,
                         use_gpu=False,
                         lang="en",
                         det_model_dir="/PaddleOCR/inference/en_PP-OCRv3_det_infer/",
                         cls_model_dir="/PaddleOCR/inference/ch_ppocr_mobile_v2.0_cls_infer/",
                         rec_model_dir="/PaddleOCR/inference/en_PP-OCRv3_rec_infer/"
                         )

# Table recognize model("图像方向分类(False)"+版面分析+表格识别)
TABLE_TEXT_REC_ENGINE = PPStructure(show_log=True,
                                    use_gpu=False,
                                    image_orientation=False,
                                    layout=True,
                                    table=True,
                                    ocr=False,
                                    mode="structure",
                                    lang="ch",
                                    det_model_dir="/PaddleOCR/inference/ch_PP-OCRv3_det_infer/",
                                    cls_model_dir="/PaddleOCR/inference/ch_ppocr_mobile_v2.0_cls_infer/",
                                    rec_model_dir="/PaddleOCR/inference/ch_PP-OCRv3_rec_infer/",
                                    # re_model_dir="/PaddleOCR/inference/re_vi_layoutxlm_xfund_infer/",
                                    # ser_model_dir="/PaddleOCR/inference/ser_vi_layoutxlm_xfund_infer/",
                                    layout_model_dir="/PaddleOCR/inference/picodet_lcnet_x1_0_fgd_layout_cdla_infer/",
                                    table_model_dir="/PaddleOCR/inference/ch_ppstructure_mobile_v2.0_SLANet_infer/"
                                    )


class OcrRecUtilBase(object):

    @staticmethod
    def cache_upload_image(filename: str = '',
                           cache_data: Union[str or bytes] = '',
                           cache_path=FILE_CACHE_PATH,
                           mode='wb'):
        """Cache images to local container cache path"""

        # Absolute filepath
        _cache_file = os.path.join(cache_path, filename)

        with open(_cache_file, mode) as fw:
            # Write bytes or str
            fw.write(cache_data)

    @classmethod
    def collect_cache_images(cls, cache_path=FILE_CACHE_PATH):
        """Collect all image file in cache_path, return sorted result."""

        # Container
        cache_images = [os.path.join(cache_path, image) for image in os.listdir(cache_path)]

        return sorted(cache_images)

    @classmethod
    def remove_recognized_cache_image(cls, file=None):

        if os.path.exists(file):
            os.remove(file)

    @classmethod
    def choose_model(cls, language: str = 'ch', rec_doc_object="ALL"):
        """ Choose model by "language and recognize doc objects"
        @language: default chinese
        @rec_doc_objet: optional objects in "SUPPORT_EXTRACT_TEXT_FROM_DOC_OBJECTS"
        """

        # Choose ocr model
        _ocr_model = None

        if rec_doc_object.upper() == "ALL":
            # Dont use layout model recognize
            if language.lower() == 'ch':
                # Use Chinese-PaddleOCR-V3-Infer-Model
                _ocr_model = OCR_CH_MODEL
            else:
                # Use English-PaddleOCR-V3-Infer-Model
                _ocr_model = OCR_EN_MODEL
        else:
            # Need use "LayoutAnalysisModel" to parse doc layout
            if rec_doc_object.upper() == "TABLE":
                # Recognize text from table-object in doc, default use chinese(If change, please change other model)
                _ocr_model = TABLE_TEXT_REC_ENGINE

            if rec_doc_object.upper() == "TEXT":
                # Recognize text from text-object in doc, default use chinese(If change, please change other model) Todo
                _ocr_model = None

            if rec_doc_object.upper() == "FIGURE":
                # Recognize text from image-object in doc, default use chinese(If change, please change other model) Todo
                _ocr_model = TABLE_TEXT_REC_ENGINE

        if _ocr_model is None:
            # Default use OCR_CH_MODEL, recognize All doc-objects.
            _ocr_model = OCR_CH_MODEL

        return _ocr_model

    @classmethod
    def extract_text_from_all_objects(cls, image_page_no: str = None, recognize_results=None):
        """Parse text from doc(image/pdf)-all-objects, dont use "LayoutAnalysisModel" to split doc-objects.
           Just use Ocr engine to recognize text from doc.
        """

        # Text container
        _text_container = list()

        if recognize_results:
            # Collect recognize text; Todo: Text lines combine
            for item in recognize_results:
                _text = item[1][0]
                _text_container.append(_text)

        # Result
        _result = {
            image_page_no: _text_container,
        }

        return _result

    @classmethod
    def extract_text_from_table_object(cls, image_page_no: str = '', recognize_results=None, rec_doc_object="TABLE"):
        """Parse text from doc(image/pdf)-table-object, use "LayoutAnalysisModel" to split doc-objects.
        then use Ocr engine to recognize text from doc-table.
        """

        # Table text container
        _tables_container = dict()

        # Todo: Need order table by bbox
        if recognize_results:
            # counter
            table_no = 0

            for item in recognize_results:
                # Get extract text from target doc object
                _doc_object = item.get("type")

                if _doc_object.upper() == rec_doc_object.upper():

                    # Accumulate table num
                    table_no += 1

                    # record: table text
                    _table_text_records = ""

                    # Extract text from table
                    _rec_res = item.get('res')

                    if _rec_res:
                        # Get html-table-text, Todo: change table format.
                        _table_text_records = _rec_res.get('html')

                    # collect table record
                    _table_no = image_page_no + "Table" + str(table_no).zfill(2)
                    _rec_tables = {_table_no: _table_text_records, }

                    # Collect recognize table records
                    _tables_container.update(_rec_tables)

        # return result
        _result = {
            image_page_no: _tables_container
        }

        return _result

    @classmethod
    def extract_text_from_text_object(cls, image_page_no: str = '', recognize_results=None, rec_doc_object="TEXT"):

        return {image_page_no: "Dont support extract text from pdf/image text-object in doc for now.", }

    @classmethod
    def extract_text_from_figure_object(cls, image_page_no: str = '', recognize_results=None, rec_doc_object="FIGURE"):

        return {image_page_no: "Dont support extract text from pdf/image figure-object in doc for now.", }

    @classmethod
    def get_recognize_text(cls, language: str = 'ch', rec_doc_object: str = 'ALL', image_cache_path=FILE_CACHE_PATH):

        # Text container
        rec_text_container = dict()

        # Timer container
        model_timer_container = dict()

        # Get model
        ocr_model = cls.choose_model(language=language, rec_doc_object=rec_doc_object)

        # Cache images file list
        cache_images_list = cls.collect_cache_images(cache_path=image_cache_path)

        # Model recognize
        for image_no, image_file in enumerate(cache_images_list, start=1):
            # keep recognize doc-objects
            rec_results = None

            # Timer: record model execute time
            model_rec_start_at = time.time()
            model_start_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")


            if rec_doc_object.upper() == "ALL":
                # Recognize text with OCR-Recognize-Engine, extract text from all doc-objects.
                rec_results = ocr_model.ocr(image_file, cls=True)

            if rec_doc_object.upper() in ['TABLE', 'TEXT', 'FIGURE']:
                # Use "Layout-Analysis" parse doc: table/text/figure-object, then ocr engine to extract text
                image_data = cv2.imread(image_file)
                rec_results = ocr_model(image_data)

            # Timer:record model execute time
            model_rec_end_at = time.time()

            # Extract text from recognize table
            image_page_no = "Page" + str(image_no).zfill(2)
            rec_text = {image_page_no: f"OCR-Model model can't recognize any results(doc_object: {rec_doc_object} in "
                                       f"{ALLOWED_EXTRACT_TEXT_FROM_DOC_OBJECTS}) from Image/Pdf.",
                        }

            if rec_results:
                # Extract text
                if rec_doc_object.upper() == "ALL":
                    # recognize text
                    rec_text = cls.extract_text_from_all_objects(image_page_no=image_page_no,
                                                                 recognize_results=rec_results)
                if rec_doc_object.upper() == 'TABLE':
                    # recognize text
                    rec_text = cls.extract_text_from_table_object(image_page_no=image_page_no,
                                                                  recognize_results=rec_results,
                                                                  rec_doc_object="TABLE")
                if rec_doc_object.upper() == 'TEXT':
                    rec_text = cls.extract_text_from_text_object(image_page_no=image_page_no,
                                                                 recognize_results=rec_results,
                                                                 rec_doc_object="TEXT")
                if rec_doc_object.upper() == 'FIGURE':
                    rec_text = cls.extract_text_from_figure_object(image_page_no=image_page_no,
                                                                   recognize_results=rec_results,
                                                                   rec_doc_object="FIGURE")

            # Timer: extract text from rec result
            extract_text_end_at = time.time()

            # Recognize success, then delete files.
            cls.remove_recognized_cache_image(file=image_file)

            # Timer: end
            image_rec_task_end_at = time.time()

            # Logger record - template
            current_app.logger.info("Stage-OcrRecImage<SubTask-%s>: recognize text from image(%s), cost: %.3f s" %
                                    (str(image_no).zfill(2), image_page_no, image_rec_task_end_at - model_rec_start_at))

            # Collect extract text
            rec_text_container.update(rec_text)

            # Collect timer in debug mode
            if current_app.debug:
                model_timer_container[image_page_no] = {
                    "ImageInfo-PageNo": image_page_no,
                    "Param01-Language": language,
                    "Param02-DocObject": rec_doc_object.upper(),
                    "Stage01-ImageRecTaskStartAt": "%s" % model_start_time,
                    "Stage02-ModelRecTaskCost": "%.3f s" % (model_rec_end_at - model_rec_start_at),
                    "Stage03-ExtractTextTaskCost": "%.3f s" % (extract_text_end_at - model_rec_end_at),
                    "Stage04-GetRecResTaskTotalCost": "%.3f s" % (extract_text_end_at - model_rec_start_at),
                    "Stage05-RemoveCacheRecImageTaskCost": "%.3f s" % (image_rec_task_end_at - extract_text_end_at),
                    "Stage06-ImageRecTaskTotalCost": "%.3f s" % (image_rec_task_end_at - model_rec_start_at)
                }

        if current_app.debug:
            # Debug model: record debug-task-timer
            rec_text_container["DebugTaskInfo"] = model_timer_container

        return rec_text_container


class ImagesUtils(OcrRecUtilBase):
    """Pictures tools """

    @staticmethod
    def cv2_to_base64(image, default_encode='utf8'):
        return base64.b64encode(image).decode(default_encode)


class PdfUtils(OcrRecUtilBase):
    """Parse pdf file object """

    @staticmethod
    def convert_pdf_to_images(pdf_file=None, pdf_stream=None, image_cache_path=None):
        """ Convert pdf (all pages) to images."""

        if pdf_file:
            _pdf = fitz.open(filename=pdf_file)
        else:
            _pdf = fitz.open(stream=pdf_stream)

        for page_no in range(_pdf.page_count):
            # convert
            page = _pdf.load_page(page_no)
            pix = page.get_pixmap()

            # Save to local path
            _page_id = str(page_no + 1).zfill(3)
            image_file = os.path.join(image_cache_path, f"ImagePage{_page_id}.png")
            pix.save(image_file)

        # Close file handle
        _pdf.close()
